package com.rabbiter.oes.util;

import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.security.SecureRandom;
import java.util.Base64;

/**
 * AES加密工具类，提供AES的加密和解密功能。
 * 支持ECB、CBC和CFB三种加密模式。
 */
public class AESUtil {

    /**
     * AES ECB加密模式字符串。
     */
    private static final String AES_ECB = "AES/ECB/PKCS5Padding";
    /**
     * AES CBC加密模式字符串。
     */
    public static final String AES_CBC = "AES/CBC/PKCS5Padding";
    /**
     * AES CFB加密模式字符串。
     */
    public static final String AES_CFB = "AES/CFB/PKCS5Padding";
    /**
     * IV向量长度。
     */
    private static final Integer IV_LENGTH = 16;
    private static final String KEY_ALGORITHM = "AES";

    /**
     * 判断字符串是否为空。
     *
     * @param str 待检查字符串。
     * @return 如果字符串为空或仅包含空格，则返回true；否则返回false。
     */
    private static boolean isEmpty(Object str) {
        return null == str || str.toString().trim().isEmpty();
    }

    /**
     * 将字符串转换为字节数组。
     *
     * @param str 待转换的字符串。
     * @return 字符串的字节数组表示，如果字符串为空，则返回null。
     */
    private static byte[] getBytes(String str) {
        if (isEmpty(str)) {
            return null;
        }
        return str.getBytes(StandardCharsets.UTF_8);
    }

    /**
     * 根据字符串生成AES密钥。
     *
     * @param keyString 密钥字符串。
     * @return AES密钥。
     * @throws IllegalArgumentException 如果密钥字符串为空，则抛出此异常。
     */
    public static SecretKeySpec getSecretKeySpec(String keyString) {
        byte[] keyBytes = getBytes(keyString);
        if (keyBytes == null) {
            throw new IllegalArgumentException("Key cannot be null or empty.");
        }
        return new SecretKeySpec(keyBytes, "AES");
    }

    /**
     * 生成安全的初始化向量（IV）。
     *
     * @return 生成的IV的Base64编码字符串。
     */
    public static String getIV() {
        SecureRandom random = new SecureRandom();
        byte[] ivBytes = new byte[IV_LENGTH];
        random.nextBytes(ivBytes);
        return Base64.getEncoder().encodeToString(ivBytes);
    }

    /**
     * 使用AES ECB模式加密字符串。
     *
     * @param text 待加密的字符串。
     * @param key  加密密钥。
     * @return 加密后的字符串，如果输入为空，则返回null。
     */
    public static String encrypt(String text, String key) {
        if (isEmpty(text) || isEmpty(key)) {
            return null;
        }
        try {
            Cipher cipher = Cipher.getInstance(AES_ECB);
            SecretKeySpec secretKeySpec = getSecretKeySpec(key);
            cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec);
            byte[] encryptedBytes = cipher.doFinal(getBytes(text));
            return Base64.getEncoder().encodeToString(encryptedBytes);
        } catch (Exception e) {
            // Consider logging this exception
            throw new RuntimeException("Encryption failed", e);
        }
    }

    /**
     * 使用AES ECB模式解密字符串。
     *
     * @param text 待解密的字符串。
     * @param key  解密密钥。
     * @return 解密后的字符串，如果输入为空，则返回null。
     */
    public static String decrypt(String text, String key) {
        if (isEmpty(text) || isEmpty(key)) {
            return null;
        }
        try {
            byte[] textBytes = Base64.getDecoder().decode(text);
            Cipher cipher = Cipher.getInstance(AES_ECB);
            SecretKeySpec secretKeySpec = getSecretKeySpec(key);
            cipher.init(Cipher.DECRYPT_MODE, secretKeySpec);
            byte[] decryptedBytes = cipher.doFinal(textBytes);
            return new String(decryptedBytes, StandardCharsets.UTF_8);
        } catch (Exception e) {
            // Consider logging this exception
            throw new RuntimeException("Decryption failed", e);
        }
    }

    /**
     * 使用指定模式和初始化向量对字符串进行AES加密。
     *
     * @param text 待加密的字符串。
     * @param key  加密密钥。
     * @param iv   初始化向量。
     * @param mode 加密模式。
     * @return 加密后的字符串，如果输入为空，则返回null。
     */
    public static String encrypt(String text, String key, String iv, String mode) {
        if (isEmpty(text) || isEmpty(key) || isEmpty(iv) || isEmpty(mode)) {
            return null;
        }
        try {
            Cipher cipher = Cipher.getInstance(mode);
            SecretKeySpec secretKeySpec = getSecretKeySpec(key);
            byte[] ivBytes = Base64.getDecoder().decode(iv);
            IvParameterSpec ivParameterSpec = new IvParameterSpec(ivBytes);
            cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec, ivParameterSpec);
            byte[] encryptedBytes = cipher.doFinal(getBytes(text));
            return Base64.getEncoder().encodeToString(encryptedBytes);
        } catch (Exception e) {
            // Consider logging this exception
            throw new RuntimeException("Encryption failed with mode: " + mode, e);
        }
    }

    /**
     * 使用指定模式和初始化向量对字符串进行AES解密。
     *
     * @param text 待解密的字符串。
     * @param key  解密密钥。
     * @param iv   初始化向量。
     * @param mode 解密模式。
     * @return 解密后的字符串，如果输入为空，则返回null。
     */
    public static String decrypt(String text, String key, String iv, String mode) {
        if (isEmpty(text) || isEmpty(key) || isEmpty(iv) || isEmpty(mode)) {
            return null;
        }
        try {
            byte[] textBytes = Base64.getDecoder().decode(text);
            Cipher cipher = Cipher.getInstance(mode);
            SecretKeySpec secretKeySpec = getSecretKeySpec(key);
            byte[] ivBytes = Base64.getDecoder().decode(iv);
            IvParameterSpec ivParameterSpec = new IvParameterSpec(ivBytes);
            cipher.init(Cipher.DECRYPT_MODE, secretKeySpec, ivParameterSpec);
            byte[] decryptedBytes = cipher.doFinal(textBytes);
            return new String(decryptedBytes, StandardCharsets.UTF_8);
        } catch (Exception e) {
            // Consider logging this exception
            throw new RuntimeException("Decryption failed with mode: " + mode, e);
        }
    }
    /**
     * 生成128位（16字节）的AES密钥。
     * @return 16字节密钥的Base64编码字符串。
     */
    public static String get16Key() {
        return generateKey(16);
    }

    /**
     * 生成192位（24字节）的AES密钥。
     * @return 24字节密钥的Base64编码字符串。
     */
    public static String get24Key() {
        return generateKey(24);
    }


    private static String generateKey(int keySizeInBytes) {
        byte[] keyBytes = new byte[keySizeInBytes];
        SecureRandom secureRandom = new SecureRandom();
        secureRandom.nextBytes(keyBytes);
        return Base64.getEncoder().encodeToString(keyBytes);
    }

    /**
     * 二次加密，防止反编译
     *
     * @param str
     * @return
     */
    public static String secondaryEncryption(String str) {
        str.replace("/","");
        str.replace("=", "");
        return str;
    }
}

